package test.utils;

import java.io.FileOutputStream;
import java.sql.Connection;
import java.sql.DriverManager;

import org.dbunit.database.DatabaseConnection;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.database.QueryDataSet;
import org.dbunit.dataset.xml.FlatXmlWriter;
import org.springframework.util.Assert;

/**
 *<p>
 * Utility for pulling data out of a database and creating a DBUnit data set
 * </p>
 * <p>
 * <b>usage: </b>DBUnitDataExtractor dbDriver dbURL dbUserId dbPassword
 * dbTableName+
 * </p>
 * 
 * @author timothystorm
 */
public class DBUnitDataExtractor {
	public static void main(String[] args) {
		Connection jdbcConnection = null;

		try {
			Assert
					.isTrue(
							(args.length >= 5),
							"usage: DBUnitDataExtractor DRIVER URL USER_ID PASS TABLE_NAME*\n"
									+ "exampe: DBUnitDataExtractor com.mysql.jdbc.Driver jdbc:mysql://localhost/fitso fitso password user food");

			String driver = args[0];
			String url = args[1];
			String userId = args[2];
			String pass = args[3];

			// Connect to the database
			Class.forName(driver);
			jdbcConnection = DriverManager.getConnection(url, userId, pass);
			IDatabaseConnection connection = new DatabaseConnection(jdbcConnection);

			// Specify the tables to retrieve the data from
			for (int i = 4; i < args.length; i++) {
				// Setup the dataSet
				QueryDataSet dataSet = new QueryDataSet(connection);
				dataSet.addTable(args[i]);

				// write dataset file
				String datasetName = args[i] + "_dataset.xml";
				FlatXmlWriter datasetWriter = new FlatXmlWriter(new FileOutputStream(datasetName));
				datasetWriter.write(dataSet);

				 //write DTD file
//				 String dtdName = args[i] + ".dtd";
//				 FlatDtdDataSet.write(dataSet, new FileOutputStream(dtdName));
			}
		}
		catch (Exception e) {
			System.out.println("Failed to create a DBUnit Dataset");
		}
	}
}
